"""
Created by Kostas Triaridis (@kostino)
in August 2023 @ ITI-CERTH
"""
import torch
import torch.nn as nn
from common.utils import SRMFilter, BayarConv2d, Fore
from models.DnCNN_noiseprint import make_net
import logging
import os


class ModalitiesExtractor(nn.Module):
    def __init__(self,
                 modals: list = ('noiseprint', 'bayar', 'srm'),
                 noiseprint_path: str = None):
        super().__init__()
        self.mod_extract = []
        if 'noiseprint' in modals:
            num_levels = 17
            out_channel = 1
            self.noiseprint = make_net(3, kernels=[3, ] * num_levels,
                                  features=[64, ] * (num_levels - 1) + [out_channel],
                                  bns=[False, ] + [True, ] * (num_levels - 2) + [False, ],
                                  acts=['relu', ] * (num_levels - 1) + ['linear', ],
                                  dilats=[1, ] * num_levels,
                                  bn_momentum=0.1, padding=1)

            if noiseprint_path:
                np_weights = noiseprint_path
                assert os.path.isfile(np_weights)
                dat = torch.load(np_weights, map_location=torch.device('cpu'))
                # dat = torch.load(np_weights)
                logging.info(f'Noiseprint++ weights: {np_weights}')
                self.noiseprint.load_state_dict(dat)

            self.noiseprint.eval()
            for param in self.noiseprint.parameters():
                param.requires_grad = False
            self.mod_extract.append(self.noiseprint)
        if 'bayar' in modals:
            self.bayar = BayarConv2d(3, 3, padding=2)
            self.mod_extract.append(self.bayar)
        if 'srm' in modals:
            self.srm = SRMFilter()
            self.mod_extract.append(self.srm)
        # if 'fore' in modals:
        #     self.fore = Fore()
        #     self.mod_extract.append(self.fore)


    def set_train(self):
        if hasattr(self, 'bayar'):
            self.bayar.train()

    def set_val(self):
        if hasattr(self, 'bayar'):
            self.bayar.eval()

    def forward(self, x) -> list:
        out = []
        for mod in self.mod_extract:
            y = mod(x)
            if y.size()[-3] == 1:
                y = torch.tile(y, (3, 1, 1))
            out.append(y)

        return out

import math
def make_net2(nplanes_in, kernels, features, bns, acts, dilats, bn_momentum=0.1, padding=1):
    depth = len(features)
    layers = []

    for i in range(depth):
        if i == 0:
            in_feats = nplanes_in
        else:
            in_feats = features[i - 1]

        # 卷积层
        conv = nn.Conv2d(in_feats, features[i], kernel_size=kernels[i],
                         dilation=dilats[i], padding=padding, bias=not bns[i])
        # 初始化
        n = kernels[i] * kernels[i] * features[i]
        conv.weight.data.normal_(0, math.sqrt(2. / n))
        layers.append(conv)

        # 批量归一化
        if bns[i]:
            bn = nn.BatchNorm2d(features[i], momentum=bn_momentum)
            n = kernels[i] * kernels[i] * features[i]
            bn.weight.data.normal_(0, math.sqrt(2. / n))
            bn.bias.data.zero_()
            layers.append(bn)

        # 激活函数
        if acts[i] == 'relu':
            layers.append(nn.ReLU(inplace=True))
        elif acts[i] == 'linear':
            pass  # 无激活

    return nn.Sequential(*layers)

class NoiseExtractor(nn.Module):
    def __init__(self,noiseprint_path: str = None):
        super().__init__()

        num_levels = 17
        out_channel = 1
        self.noiseprint = make_net2(3, kernels=[3] * num_levels,
                              features=[64] * (num_levels - 1) + [out_channel],
                              bns=[False] + [True] * (num_levels - 2) + [False],
                              acts=['relu'] * (num_levels - 1) + ['linear'],
                              dilats=[1] * num_levels,
                              bn_momentum=0.1, padding=1)

        np_weights = noiseprint_path
        assert os.path.isfile(np_weights)
        dat = torch.load(np_weights, map_location=torch.device('cpu'))
        # dat = torch.load(np_weights)
        logging.info(f'Noiseprint++ weights: {np_weights}')
        self.noiseprint.load_state_dict(dat)

        self.noiseprint.eval()
        for param in self.noiseprint.parameters():
            param.requires_grad = False



    def forward(self, x):
        y = self.noiseprint(x)
        if y.size()[-3] == 1:
            y = torch.tile(y, (3, 1, 1))
        return y


if __name__ == '__main__':
    modal_ext = ModalitiesExtractor(['noiseprint', 'bayar', 'srm','fore'], '../pretrained/noiseprint/np++.pth')
    # modal_ext = NoiseExtractor('/home/wc/disk1/MMFusion/dncnn_epoch_50.pth')

    from PIL import Image
    from torchvision.transforms.functional import to_tensor, to_pil_image
    # img = Image.open('/home/wc/disk1/datasets/COLUMBIA/Columbia/4cam_splc/canong3_kodakdcs330_sub_06.tif').convert('RGB')
    # img = Image.open('/home/wc/disk1/datasets/COVER/tampered/6t.tif').convert('RGB')
    img = Image.open('/home/wc/disk1/DocTamper/DocTamperV1/DocTamperV1-TestingSet/img/0.jpg').convert('RGB')
    inp = to_tensor(img).unsqueeze(0)

    out = modal_ext(inp)

    import matplotlib.pyplot as plt
    import numpy as np
    fig, ax = plt.subplots(2, 2, figsize=(28,28))
    plt.figure(figsize=(8, 6),  dpi=200)
    ax[0,0].imshow(img)
    # ax[0,0].set_title('Image')
    # plt.imshow(img)
    noiseprint = out[0][:, 0].squeeze().numpy()
    ax[0,1].imshow(noiseprint[16:-16:4, 16:-16:4], cmap='gray')
    # plt.imshow(noiseprint[16:-16:4, 16:-16:4], cmap='gray')  # 单独
    # ax[0,1].set_title('NoisePrint++')
    #
    # bayar = to_tensor(to_pil_image(out[1].squeeze())).permute(1, 2, 0).numpy()
    # ax[1,0].imshow(bayar)
    # # plt.imshow(bayar)  # 单独
    # # ax[1,0].set_title('Bayar')
    # # #
    # srm = to_tensor(to_pil_image(out[2].squeeze())).permute(1, 2, 0).numpy()
    # ax[1,1].imshow(srm)
    # # plt.imshow(srm)  # 单独
    # # ax[1,1].set_title('SRM')
    # # fore = to_tensor(to_pil_image(out[3][0:1,:,:].squeeze())).permute(1, 2, 0).numpy()
    # # ax[3].imshow(fore.squeeze(), cmap='gray')
    # ax[3].set_title('FORE')

    plt.show()

    # plt.close()
